#pragma once

#include <algorithm>
#include <initializer_list>
#include <numeric>
#include <vector>
#include <optional>

#include "TxDNN/txdnn.h"


struct txdnnTensorStruct {};

namespace txdnn {

struct TensorDescriptor : public txdnnTensorStruct {
public:
  TensorDescriptor() : is_packed_(true) {}
  TensorDescriptor(txdnnDataType_t t) : is_packed_(true), type_(t) {}

  TensorDescriptor(txdnnDataType_t t, const std::initializer_list<int>& dims)
    : TensorDescriptor(t, std::vector<int>(dims)) {}
  
  TensorDescriptor(txdnnDataType_t t, const std::vector<int>& dims);

  TensorDescriptor(txdnnDataType_t t, const std::initializer_list<std::size_t>& dims)
    : TensorDescriptor(t, std::vector<std::size_t>(dims)) {}
  
  TensorDescriptor(txdnnDataType_t t, const std::vector<std::size_t>& dims);

  TensorDescriptor(txdnnDataType_t t, std::vector<std::size_t>&& dims);

  TensorDescriptor(txdnnDataType_t t, txdnnBackendLayoutType_t layout, const std::vector<int>& dims);

  TensorDescriptor(txdnnDataType_t t, txdnnBackendLayoutType_t layout, const std::initializer_list<std::size_t>& dims)
    : TensorDescriptor(t, layout, std::vector<std::size_t>(dims)) {}

  TensorDescriptor(txdnnDataType_t t, txdnnBackendLayoutType_t layout, const std::vector<std::size_t>& dims)
    : TensorDescriptor(t, layout, dims, {}, false) {}

  TensorDescriptor(txdnnDataType_t t, txdnnBackendLayoutType_t layout, std::vector<std::size_t>&& dims)
    : TensorDescriptor(t, layout, std::move(dims), {}, false) {}

  TensorDescriptor(txdnnDataType_t t, const std::vector<int>& dims, const std::vector<int>& strides);

  TensorDescriptor(txdnnDataType_t t, const std::initializer_list<std::size_t>& dims,
                   const std::initializer_list<std::size_t>& strides)
    : TensorDescriptor(t, std::vector<std::size_t>(dims), std::vector<std::size_t>(strides)) {}

  TensorDescriptor(txdnnDataType_t t, const std::vector<std::size_t>& dims, const std::vector<std::size_t>& strides)
    : TensorDescriptor(t, std::nullopt, dims, strides, true) {}

  TensorDescriptor(txdnnDataType_t t, std::vector<std::size_t>&& dims, std::vector<std::size_t>&& strides)
    : TensorDescriptor(t, std::nullopt, std::move(dims), std::move(strides), true) {}

  TensorDescriptor(txdnnDataType_t t, txdnnBackendLayoutType_t layout, const std::vector<std::size_t>& dims,
                   const std::vector<std::size_t>& strides)
    : TensorDescriptor(t, layout, dims, strides, true) {}

  TensorDescriptor(txdnnDataType_t t, txdnnBackendLayoutType_t layout, std::vector<std::size_t>&& dims,
                   std::vector<std::size_t>&& strides)
    : TensorDescriptor(t, layout, std::move(dims), std::move(strides), true) {}

  static TensorDescriptor makeDescriptor(txdnnDataType_t t, const int* dims, int size);
  static TensorDescriptor makeDescriptor(txdnnDataType_t t, const std::size_t* dims, int size);
  static TensorDescriptor makeDescriptor(txdnnDataType_t t, txdnnBackendLayoutType_t layout, const int* dims, int size);
  static TensorDescriptor makeDescriptor(txdnnDataType_t t, txdnnBackendLayoutType_t layout, const std::size_t* dims,
                                         int size);
  static TensorDescriptor makeDescriptor(txdnnDataType_t t, const int* dims, const int* strides, int size);
  static TensorDescriptor makeDescriptor(txdnnDataType_t t, const std::size_t* dims, const std::size_t* strides,
                                         int size);

  txdnnDataType_t getDataType() const { return type_; }

  const std::vector<std::size_t>& getDims() const { return dims_; }
  const std::vector<std::size_t>& getStrides() const { return strides_; }

  unsigned getNumDims() const { return dims_.size(); }
  std::size_t getElementSize() const {
    return std::accumulate(dims_.begin(), dims_.end(), vector_length_, std::multiplies<std::size_t>());
  }

private:
  TensorDescriptor(txdnnDataType_t t,
                   const std::optional<txdnnBackendLayoutType_t>& layout,
                   const std::vector<std::size_t>& dims,
                   const std::vector<std::size_t>& strides,
                   bool use_strides)
    : type_(t), layout_(layout), dims_(dims), strides_(use_strides ? strides : std::vector<std::size_t>()) {
    this->checkArgsAndInit(use_strides);
  }

  TensorDescriptor(txdnnDataType_t t, 
                   const std::optional<txdnnBackendLayoutType_t>& layout,
                   const std::vector<std::size_t>&& dims,
                   const std::vector<std::size_t>&& strides,
                   bool use_strides)
    : type_(t), layout_(layout), dims_(dims), strides_(use_strides ? std::move(strides) : std::vector<std::size_t>()) {
    this->checkArgsAndInit(use_strides);
  }

  void checkArgsAndInit(bool use_strides);

private:
  std::vector<std::size_t> dims_;
  std::vector<std::size_t> strides_;

  bool is_packed_ = true;
  std::size_t vector_length_ = 1;

  txdnnDataType_t type_ = TXDNN_DATA_FLOAT;
  std::optional<txdnnDataType_t> cast_type_;
  std::optional<txdnnBackendLayoutType_t> layout_;
};

} // namespace txdnn

TXDNN_OBJECT_DEFINE(txdnnTensorStruct, txdnn::TensorDescriptor);